# 仅加密场景
import math
import random
import struct
import threading
from contextlib import nullcontext
from pathlib import Path

from rich.progress import Progress

from storage.blobstore.crypto import encrypt_blob_aes128_ctr, encrypt_blob, decrypt_blob
from storage.blobstore.fileblob import EncryptionAlgorithm
from storage.utils import mmap_view, common_progress

MAGIC_HEADER = b'ZDEC'
FILE_HEADER_STRUCT_FMT = '>4sHLQ'
BLCK_HEADER_STRUCT_FMT = '>16sL'

FILE_HEADER_SIZE = struct.calcsize(FILE_HEADER_STRUCT_FMT)
BLCK_HEADER_SIZE = struct.calcsize(BLCK_HEADER_STRUCT_FMT)


def _calc_final_size(total: int, block_size: int, alg: EncryptionAlgorithm) -> int:
    """
    计算加密后文件的大小

    此函数根据源文件的总大小和分块的大小，计算加密后文件的大小。
    它考虑了文件头、分块头和分块内容的大小。

    计算公式 = 文件头大小 + 总块数(向上取整) * (块大小+块头大小+加密块大小(如果是block cipher的话))

    :param total: 源文件的总大小
    :param block_size: 分块的大小
    :return: 加密后文件的大小
    """
    return FILE_HEADER_SIZE + math.ceil(total / block_size) * (
        BLCK_HEADER_SIZE + block_size + alg.block_size if alg.is_block_cipher else 0)


def _calc_encrypted_offset(block_index: int, block_size: int, alg: EncryptionAlgorithm) -> int:
    """
    计算分块在加密后文件中的偏移量

    此函数根据分块的索引和分块的大小，计算分块在加密后文件中的偏移量。
    它考虑了文件头和分块头的大小。

    :param block_index: 分块的索引
    :param block_size: 分块的大小
    :return: 分块在加密后文件中的偏移量
    """
    return FILE_HEADER_SIZE + block_index * (
        BLCK_HEADER_SIZE + block_size + alg.block_size if alg.is_block_cipher else 0)


def _decrypt_range(src: memoryview, dest: memoryview, alg: EncryptionAlgorithm, key: bytes, block_size: int,
                   show_progress: bool = False, progress: Progress = None, discard_output: bool = False):
    """
    将源数据范围解密到目标。

    此函数从源内存视图中按块读取数据，使用指定的算法和密钥进行解密，并将解密后的数据写入目标内存视图。它还支持显示解密过程的进度。

    :param src: 要解密的源数据的内存视图。
    :param dest: 解密数据将写入的目标内存视图。
    :param alg: 要使用的加密算法。
    :param key: 解密密钥。
    :param block_size: 解密的数据块大小。
    :param show_progress: 是否显示解密进度。
    :param progress: 进度跟踪对象。
    :param discard_output: 如果为 True，则解密数据不会写入目标。
    """
    task = None
    if show_progress:
        task = progress.add_task(f'decryption thread', total=len(src))

    encrypted_block_size = BLCK_HEADER_SIZE + block_size + alg.block_size if alg.is_block_cipher else 0
    with src, dest:
        for block_index in range(len(src) // encrypted_block_size):
            src_from = block_index * encrypted_block_size
            src_to = src_from + encrypted_block_size
            if src_to <= src_from:
                break

            with src[src_from:src_from + BLCK_HEADER_SIZE] as data:
                nonce_iv, real_size = struct.unpack(BLCK_HEADER_STRUCT_FMT, data)
            with src[src_from + BLCK_HEADER_SIZE:src_to] as data:
                decrypted = decrypt_blob(alg, key, nonce_iv, data)

            if not discard_output:
                dest[block_index * block_size:block_index * block_size + real_size] = decrypted[:real_size]

            if show_progress:
                progress.update(task, advance=encrypted_block_size)


def _encrypt_range(src: memoryview, dest: memoryview, alg: EncryptionAlgorithm, key: bytes, block_size: int,
                   show_progress: bool = False, progress: Progress = None, discard_output: bool = False):
    """
    将源数据范围加密到目标。

    该函数从源内存视图中按块读取数据，使用指定的算法和密钥进行加密，并将加密后的数据写入目标内存视图。它还支持显示加密过程的进度。

    :param src: 要加密的源数据的内存视图。
    :param dest: 加密数据将写入的目标内存视图。
    :param alg: 要使用的加密算法。
    :param key: 加密密钥。
    :param block_size: 加密的数据块大小。
    :param show_progress: 是否显示加密进度。
    :param progress: 进度跟踪对象。
    :param discard_output: 如果为 True，则加密数据不会写入目标。
    """
    task = None
    if show_progress:
        task = progress.add_task(f'encryption thread', total=len(src))

    encrypted_block_size = BLCK_HEADER_SIZE + block_size + alg.block_size if alg.is_block_cipher else 0
    with src, dest:
        # 按指定大小的块迭代源数据
        for block_index in range(len(src) // block_size):
            src_from = block_index * block_size
            # 计算当前块的结束位置
            src_to = src_from + block_size
            # 如果结束位置不大于起始位置，退出循环
            if src_to <= src_from:
                break

            # 生成用于 AES 加密的随机初始化向量 (nonce_iv)
            nonce_iv = random.randbytes(alg.nonce_iv_size)

            # 加密当前数据块
            if src_to - src_from < block_size:
                # 不足 block_size 的块填充到 block_size 大小
                with src[src_from:src_to] as data:
                    data = bytes(data) + b'\0' * (block_size - len(data))
                    ciphertext = encrypt_blob_aes128_ctr(key, nonce_iv, data)
                    assert len(ciphertext) == block_size + alg.block_size if alg.is_block_cipher else 0
            else:
                with src[src_from:src_to] as data:
                    ciphertext = encrypt_blob(alg, key, nonce_iv, data)
                    assert len(ciphertext) == block_size + alg.block_size if alg.is_block_cipher else 0

            # 如果输出不被丢弃，将加密数据和头部写入目标
            if not discard_output:
                # 计算目标缓冲区中的对应位置
                dest_from = block_index * encrypted_block_size
                dest_to = dest_from + encrypted_block_size
                # 如果目标缓冲区太小，引发错误
                if dest_to > len(dest):
                    raise ValueError(f'dest buffer too small, requires {dest_to} bytes, got {len(dest)}')

                # 打包加密头部和加密数据，然后写入目标缓冲区
                packed = struct.pack(BLCK_HEADER_STRUCT_FMT, nonce_iv, src_to - src_from) + ciphertext
                dest[dest_from:dest_to] = packed

            # 如果启用了进度显示，更新进度
            if show_progress:
                progress.update(task, advance=src_to - src_from)


def decrypt_file(src: Path, dst: Path, key: bytes, /, *,
                 parallel: int = 1, show_progress: bool = False, discard_output: bool = False):
    progress = common_progress(show_progress)
    with mmap_view(src, readonly=True) as src_view, progress:
        # 校验魔术头
        with src_view[:FILE_HEADER_SIZE] as data:
            magic, alg, block_size, total = struct.unpack(FILE_HEADER_STRUCT_FMT, data)
        if magic != MAGIC_HEADER:
            raise ValueError('invalid file magic')

        # 验证加密算法枚举
        alg = EncryptionAlgorithm(alg)
        # 确定加密后块大小
        encrypted_block_size = BLCK_HEADER_SIZE + block_size + (alg.block_size if alg.is_block_cipher else 0)
        # 确定全文件的 block 数量
        n_blocks = (len(src_view) - FILE_HEADER_SIZE) // encrypted_block_size
        # n_blocks 分配给 parallel 个线程，剩余的 blocks 分配给最后一个线程
        thread_n_blocks = n_blocks // parallel
        # 待启动线程
        threads = []

        if not discard_output:
            dst_ctx = mmap_view(dst, readonly=False, truncate=total)
        else:
            dst_ctx = nullcontext()

        with dst_ctx as dst_view:
            for i in range(parallel):
                blck_from = i * thread_n_blocks
                blck_to = (i + 1) * thread_n_blocks
                if i == parallel - 1:
                    blck_to += n_blocks % parallel

                # 计算源缓冲区的起始和结束位置
                src_from, src_to = _calc_encrypted_offset(blck_from, block_size, alg), _calc_encrypted_offset(blck_to,
                                                                                                              block_size,
                                                                                                              alg)
                dst_from, dst_to = blck_from * block_size, blck_to * block_size

                input_view = src_view[src_from:src_to]
                if not discard_output:
                    output_view = dst_view[dst_from:dst_to]
                else:
                    output_view = nullcontext()

                args = (input_view, output_view, alg, key, block_size)
                kwargs = dict(show_progress=show_progress, progress=progress, discard_output=discard_output)
                thread = threading.Thread(target=_decrypt_range, args=args, kwargs=kwargs)
                threads.append(thread)

            # 启动所有线程进行解压缩，并等待所有线程完成
            for thread in threads:
                thread.start()

            for thread in threads:
                thread.join()


def encrypt_file(src: Path, dst: Path, block_size: int, alg: EncryptionAlgorithm, key: bytes, /, *,
                 parallel: int = 1, show_progress: bool = False, discard_output: bool = False):
    progress = common_progress(show_progress)
    if not discard_output:
        dst_ctx = mmap_view(dst, readonly=False, truncate=_calc_final_size(src.stat().st_size, block_size, alg))
    else:
        dst_ctx = nullcontext()

    with mmap_view(src, readonly=True) as src_view, dst_ctx as dst_view, progress:
        # 确定全文件的 block 数量
        n_blocks = math.ceil(len(src_view) // block_size)
        # n_blocks 分配给 parallel 个线程，剩余的 blocks 分配给最后一个线程
        thread_n_blocks = n_blocks // parallel
        # 待启动线程
        threads = []

        # 构造线程
        for i in range(parallel):
            # 开始块和结束块序号
            blck_from = i * thread_n_blocks
            blck_to = (i + 1) * thread_n_blocks
            # 最后一个线程处理不能均分的剩余块
            if i == parallel - 1:
                blck_to += n_blocks % parallel

            # 计算分块的范围
            src_from, src_to = blck_from * block_size, blck_to * block_size
            dst_from, dst_to = _calc_encrypted_offset(blck_from, block_size, alg), _calc_encrypted_offset(blck_to,
                                                                                                          block_size,
                                                                                                          alg)

            # 创建输入输出视图
            input_view = src_view[src_from:src_to]
            if not discard_output:
                output_view = dst_view[dst_from:dst_to]
            else:
                output_view = nullcontext()

            # 提交加密任务
            args = (input_view, output_view, alg, key, block_size)
            kwargs = dict(show_progress=show_progress, progress=progress, discard_output=discard_output)
            thread = threading.Thread(target=_encrypt_range, args=args, kwargs=kwargs)
            threads.append(thread)

        # 等待所有任务完成
        for thread in threads:
            thread.start()

        for thread in threads:
            thread.join()

        # 写入文件头
        if not discard_output:
            with dst_view[0:FILE_HEADER_SIZE] as header:
                header[:] = struct.pack(FILE_HEADER_STRUCT_FMT, MAGIC_HEADER, alg, block_size, len(src_view))
